#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include "params.h"
#include "sign.h"
#include "packing.h"
#include "polyvec.h"
#include "poly.h"
#include "randombytes.h"
#include "symmetric.h"
#include "fips202.h"



/*************************************************
* Name:        crypto_sign_keypair
*
* Description: Generates public and private key.
*
* Arguments:   - uint8_t *pk: pointer to output public key (allocated
*                             array of CRYPTO_PUBLICKEYBYTES bytes)
*              - uint8_t *sk: pointer to output private key (allocated
*                             array of CRYPTO_SECRETKEYBYTES bytes)
*
* Returns 0 (success)
**************************************************/
//uint8_t fixedseed[32] = {0x06,0x10,0x67,0x8f,0xf4,0xdc,0x31,0x28,0xe1,0x61,0x9f,0x91,0x5d,0xc1,0x92,0xc2,0x20,0xf8,0xfa,0xd9,0x4d,0xa1,0x94,0x3b,0x90,0xaa,0xec,0x40,0x16,0x83,0xa4,0x92};
  //uint8_t fixedseed[32] = {0xef,0x99,0x22,0x4a,0x03,0xa8,0x5a,0x46,0xef,0x11,0x54,0x74,0xec,0x5b,0x5d,0x62,0x0d,0xa6,0x79,0x5d,0x6e,0xfc,0xca,0x4c,0x91,0x35,0xd1,0x99,0x58,0xa9,0xde,0x62};
  //uint8_t fixedseed[32] = {0x9f,0x52,0xaf,0x92,0xca,0x16,0x5f,0xdc,0x38,0x78,0x8f,0x2b,0x59,0xba,0x02,0xe0,0x1c,0x82,0x81,0xff,0x7c,0x1e,0x60,0x50,0x46,0x88,0x04,0x3a,0x5f,0xe8,0x14,0xb0};
  //uint8_t fixedseed[32] = {0x61,0x0a,0xfb,0x64,0xbe,0x8c,0xc1,0xdf,0x28,0x8c,0xfb,0x01,0x6e,0xe2,0xf4,0x4c,0x6c,0x07,0x11,0x3d,0xe7,0xf6,0xfe,0xe0,0x71,0xfe,0x0c,0x3f,0xe3,0x1c,0x62,0x15};
  //uint8_t fixedseed[32] = {0x3e,0x80,0x9e,0xc8,0xdd,0x0f,0xec,0x0d,0x91,0x1a,0x4e,0x3f,0xac,0x20,0xf7,0x0f,0xbb,0x12,0x8c,0x5d,0xe9,0x4d,0xc7,0x18,0x4c,0xa7,0x31,0x0a,0xe9,0x15,0x7a,0x98};
  //uint8_t fixedseed[32] = {0x24,0x78,0xf7,0xd3,0xde,0x60,0x41,0xe7,0xe5,0xcd,0x11,0xc5,0xe2,0xef,0x48,0x3d,0x1a,0xa6,0x21,0x8e,0xb1,0x26,0x44,0x40,0x91,0x53,0x5f,0x6a,0xe5,0x32,0xfa,0x73};
  //uint8_t fixedseed[32] = {0xbb,0xd1,0x54,0x19,0x06,0x0c,0xab,0x04,0x04,0xe4,0xb0,0x26,0xb5,0xef,0x97,0x5e,0x32,0x73,0xc0,0x87,0x41,0xc9,0xb2,0xdf,0x6a,0xdc,0x7d,0xd8,0x5e,0xb5,0x39,0x23};
  //uint8_t fixedseed[32] = {0x4f,0xc6,0x0f,0x1d,0x6f,0x8c,0x6f,0xad,0x2a,0x7f,0xd5,0x0d,0x73,0xb5,0xdd,0x3f,0xd9,0x3a,0x95,0x0e,0x33,0x7d,0x66,0xed,0x04,0x26,0x83,0x80,0x9f,0xb6,0x6a,0x1d};
//uint8_t fixedseed[32] = {0x7d,0xdb,0x25,0xcb,0xc1,0x03,0xda,0x5c,0xd0,0x8c,0xd2,0xa4,0x0d,0xb1,0x59,0x82,0xe1,0xcb,0x13,0xee,0x28,0x72,0x7b,0xb8,0xe5,0x49,0x88,0x6f,0xbb,0xb2,0x52,0x36};
//uint8_t fixedseed[32] = {0x6c,0x39,0x1e,0xaa,0x01,0xa5,0xb5,0x08,0x12,0x4b,0x43,0x0d,0xff,0x5b,0x74,0xb3,0xdf,0xfc,0xd3,0x62,0x15,0x45,0x97,0x22,0xe3,0xfd,0x04,0xd6,0xb8,0x3a,0x87,0xa7};

//uint8_t fixedseed[32] = {0x82,0xd0,0xd3,0xa2,0xfb,0x36,0x98,0x11,0x93,0xae,0x30,0x47,0x2e,0x73,0x4f,0x4b,0x41,0x3c,0xf3,0x06,0xf1,0xeb,0x72,0x12,0x80,0x4e,0xea,0x5a,0xb9,0x79,0x13,0x30};
//uint8_t fixedseed[32] = {0xe9,0x0a,0x5a,0xbd,0xe2,0x1f,0xe3,0x94,0x4d,0xdc,0xbb,0xd1,0x6c,0x26,0x2a,0x7c,0x5a,0x8c,0xdc,0x49,0x0b,0x5d,0x4f,0xc3,0x99,0x02,0x1b,0x77,0x1d,0x9e,0xfa,0x74};
//uint8_t fixedseed[32] = {0xbb,0xa9,0x15,0xa6,0xb5,0x4a,0x52,0x72,0x90,0x31,0x5f,0xb5,0x51,0x92,0xc8,0x7f,0x17,0xb0,0x9e,0x22,0xa3,0x40,0x7c,0x2c,0x8f,0x6e,0x79,0xd1,0x62,0xac,0xf6,0xc1};
//uint8_t fixedseed[32] = {0x84,0x92,0x99,0xff,0x59,0xe9,0xf9,0x38,0x68,0x19,0x62,0x4d,0xc2,0x53,0x73,0xf6,0xa7,0xe6,0x3b,0x1a,0x93,0xab,0x14,0x56,0x7b,0x71,0x87,0x15,0xd7,0xee,0x51,0x11};
uint8_t fixedseed[32] = {0x16,0xc2,0xd8,0xa1,0x29,0x89,0x96,0x87,0x94,0x86,0xca,0x38,0x38,0x55,0x91,0xd7,0xe3,0x62,0x24,0xbb,0x72,0x8a,0x42,0x14,0x1a,0x54,0xc5,0x6e,0x07,0xb5,0x8e,0xde};
uint8_t fixedseed[32]={0x33, 0x79, 0xF7, 0xBC, 0x63, 0xB3, 0x70, 0x83, 0x60, 0x04, 0xF8, 0x3A, 0x84, 0x1C, 0x4B, 0xFD, 0xC2, 0xE2, 0x33, 0x8B, 0x37, 0xA7, 0x85, 0x62, 0x33, 0x23, 0xC9, 0x32, 0xAB, 0x8D, 0x8C, 0xA7};
int crypto_sign_keypair(uint8_t *pk, uint8_t *sk) {
  uint8_t seedbuf[3*SEEDBYTES];
  uint8_t tr[CRHBYTES];
  const uint8_t *rho, *rhoprime, *key;
  polyvecl mat[K];
  polyvecl s1, s1hat;
  polyveck s2, t1, t0;

  for (int i = 0; i < SEEDBYTES; i++){
      seedbuf[i] = fixedseed[i];
  }
  /* Get randomness for rho, rhoprime and key */
  //randombytes(seedbuf, SEEDBYTES);
  
  shake256(seedbuf, 3*SEEDBYTES, seedbuf, SEEDBYTES);
  

  rho = seedbuf;
  rhoprime = seedbuf + SEEDBYTES;
  key = seedbuf + 2*SEEDBYTES;

  /* Expand matrix */
  polyvec_matrix_expand(mat, rho);

  /* Sample short vectors s1 and s2 */
  polyvecl_uniform_eta(&s1, rhoprime, 0);
  polyveck_uniform_eta(&s2, rhoprime, L);

  /* Matrix-vector multiplication */
  s1hat = s1;
  polyvecl_ntt(&s1hat);
  polyvec_matrix_pointwise_montgomery(&t1, mat, &s1hat);
  polyveck_reduce(&t1);
  polyveck_invntt_tomont(&t1);

  /* Add error vector s2 */
  polyveck_add(&t1, &t1, &s2);

  /* Extract t1 and write public key */
  polyveck_caddq(&t1);
  polyveck_power2round(&t1, &t0, &t1);
  pack_pk(pk, rho, &t1);

  /* Compute CRH(rho, t1) and write secret key */
  crh(tr, pk, CRYPTO_PUBLICKEYBYTES);
  pack_sk(sk, rho, tr, key, &t0, &s1, &s2);

  return 0;
}


/*************************************************
* Name:        crypto_sign_signature
*
* Description: Computes signature.
*
* Arguments:   - uint8_t *sig:   pointer to output signature (of length CRYPTO_BYTES)
*              - size_t *siglen: pointer to output length of signature
*              - uint8_t *m:     pointer to message to be signed
*              - size_t mlen:    length of message
*              - uint8_t *sk:    pointer to bit-packed secret key
*
* Returns 0 (success)
**************************************************/
int crypto_sign_signature(uint8_t *sig,
                          size_t *siglen,
                          const uint8_t *m,
                          size_t mlen,
                          const uint8_t *sk)
{
  unsigned int n;
  uint8_t seedbuf[2*SEEDBYTES + 3*CRHBYTES];
  uint8_t *rho, *tr, *key, *mu, *rhoprime;
  uint16_t nonce = 0;
  polyvecl mat[K], s1, y, z;
  polyveck t0, s2, w1, w0, h;
  poly cp;
  keccak_state state;

  rho = seedbuf;
  tr = rho + SEEDBYTES;
  key = tr + CRHBYTES;
  mu = key + SEEDBYTES;
  rhoprime = mu + CRHBYTES;
  unpack_sk(rho, tr, key, &t0, &s1, &s2, sk);

  /* Compute CRH(tr, msg) */
  shake256_init(&state);
  shake256_absorb(&state, tr, CRHBYTES);
  shake256_absorb(&state, m, mlen);
  shake256_finalize(&state);
  shake256_squeeze(mu, CRHBYTES, &state);

#ifdef DILITHIUM_RANDOMIZED_SIGNING
//no here
  randombytes(rhoprime, CRHBYTES);
#else
//printf("here");
//by shake256
  crh(rhoprime, key, SEEDBYTES + CRHBYTES);
#endif

  /* Expand matrix and transform vectors */
  polyvec_matrix_expand(mat, rho);
  polyvecl_ntt(&s1);
  polyveck_ntt(&s2);
  polyveck_ntt(&t0);

rej:
  /* Sample intermediate vector y */
  //here!!!!!!!!!!!!!!!!!!!!!!!!!
  
  polyvecl_uniform_gamma1(&y, rhoprime, nonce++);

  //here!!!!!!!!!!!!!!!!!!!!!!!!
  z = y;
  polyvecl_ntt(&z);

  /* Matrix-vector multiplication */
  polyvec_matrix_pointwise_montgomery(&w1, mat, &z);
  polyveck_reduce(&w1);
  polyveck_invntt_tomont(&w1);

  /* Decompose w and call the random oracle */
  polyveck_caddq(&w1);
  polyveck_decompose(&w1, &w0, &w1);
  polyveck_pack_w1(sig, &w1);

  shake256_init(&state);
  shake256_absorb(&state, mu, CRHBYTES);
  shake256_absorb(&state, sig, K*POLYW1_PACKEDBYTES);
  shake256_finalize(&state);
  shake256_squeeze(sig, SEEDBYTES, &state);
  poly_challenge(&cp, sig);
  poly_ntt(&cp);

  /* Compute z, reject if it reveals secret */
  polyvecl_pointwise_poly_montgomery(&z, &cp, &s1);
  polyvecl_invntt_tomont(&z);
  polyvecl_add(&z, &z, &y);
  polyvecl_reduce(&z);
  if(polyvecl_chknorm(&z, GAMMA1 - BETA)){
    //printf("WRONG!!!-------");
    //printf("/n");
   // return 1;
    goto rej;
    }

  /* Check that subtracting cs2 does not change high bits of w and low bits
   * do not reveal secret information */
  polyveck_pointwise_poly_montgomery(&h, &cp, &s2);
  polyveck_invntt_tomont(&h);
  polyveck_sub(&w0, &w0, &h);
  polyveck_reduce(&w0);
  if(polyveck_chknorm(&w0, GAMMA2 - BETA)){
      //printf("WRONG!!!-------");
    //printf("/n");
   // return 1;
    goto rej;
    }

  /* Compute hints for w1 */
  polyveck_pointwise_poly_montgomery(&h, &cp, &t0);
  polyveck_invntt_tomont(&h);
  polyveck_reduce(&h);
  if(polyveck_chknorm(&h, GAMMA2)){
      //printf("WRONG!!!-------");
    //printf("/n");
   // return 1;
    goto rej;
    }

  polyveck_add(&w0, &w0, &h);
  polyveck_caddq(&w0);
  n = polyveck_make_hint(&h, &w0, &w1);
  if(n > OMEGA){
      //printf("WRONG!!!-------");
    //printf("/n");
   // return 1;
    goto rej;
    }

  /* Write signature */
  pack_sig(sig, sig, &z, &h);
  *siglen = CRYPTO_BYTES;
  
  

  /* Get randomness for rho, rhoprime and key */
  //randombytes(seedbuf, SEEDBYTES);
  

  
  uint8_t mark[4] = {1,1,1,1};
  uint8_t index = 0;
  for(int l = 0 ; l < 4 ; l ++){
  	for(int i = 0; i < 256; i++){
  		if(y.vec[l].coeffs[i] == 0) {
	  		mark[l] = 0;
  			break;
  		}
  	}
  	
  }
  
  if(mark[0] == 1 && mark[1] == 1 && mark[2] == 1 && mark[3]==1){
  	//printf("NO_ZERO_WRONG");
  	//printf("--");
  	return -1;
  }else{
      
      for(index = 0; index < 4; index ++){
      	if(mark[index] == 0) break;
          
      }
  
  }
  extern unsigned long long global_counter;
  printf("counter=%llu\n", global_counter);
  printf("mark=");
  for(int l = 0 ; l < 4 ; l ++){
  	printf("%d ", mark[l]);
  }
  printf("\n");
  
  printf("rhoprime=");
  for (int i = 0; i < CRHBYTES; i++){
      printf("%02x", rhoprime[i]);
  }
  printf("\n");
  
  printf("nonce=");

  printf("%04x", nonce);
  printf("\n");
  
  
  
  printf("mark_index=%d", index);
  printf("\n");

  printf("r_f4bytes[%d]=", index);
  for(int j = 0;j<16; j++ ){
  	printf("%02x", *((uint8_t*)(y.vec[index].coeffs)+j));
  }
  printf("\n");
  printf("r_l4bytes[%d]=", index);
  for(int j =0;j < 16; j++){
  	printf("%02x", *((uint8_t*)(y.vec[index].coeffs)+256 * 4 - 16 +j));
  }
  printf("\n");
  printf("r_allbytes[%d]=", index);
   for(int j = 0;j<256*4; j++ ){
  	printf("%02x", *((uint8_t*)(y.vec[index].coeffs)+j));
  }
  printf("\n");
  

  

  	  
  
  
  
  return 0;
}

/*************************************************
* Name:        crypto_sign
*
* Description: Compute signed message.
*
* Arguments:   - uint8_t *sm: pointer to output signed message (allocated
*                             array with CRYPTO_BYTES + mlen bytes),
*                             can be equal to m
*              - size_t *smlen: pointer to output length of signed
*                               message
*              - const uint8_t *m: pointer to message to be signed
*              - size_t mlen: length of message
*              - const uint8_t *sk: pointer to bit-packed secret key
*
* Returns 0 (success)
**************************************************/
int crypto_sign(uint8_t *sm,
                size_t *smlen,
                const uint8_t *m,
                size_t mlen,
                const uint8_t *sk)
{
  size_t i;

  for(i = 0; i < mlen; ++i)
    sm[CRYPTO_BYTES + mlen - 1 - i] = m[mlen - 1 - i];
  int ret = crypto_sign_signature(sm, smlen, sm + CRYPTO_BYTES, mlen, sk);
  *smlen += mlen;
  return ret;
}

/*************************************************
* Name:        crypto_sign_verify
*
* Description: Verifies signature.
*
* Arguments:   - uint8_t *m: pointer to input signature
*              - size_t siglen: length of signature
*              - const uint8_t *m: pointer to message
*              - size_t mlen: length of message
*              - const uint8_t *pk: pointer to bit-packed public key
*
* Returns 0 if signature could be verified correctly and -1 otherwise
**************************************************/
int crypto_sign_verify(const uint8_t *sig,
                       size_t siglen,
                       const uint8_t *m,
                       size_t mlen,
                       const uint8_t *pk)
{
  unsigned int i;
  uint8_t buf[K*POLYW1_PACKEDBYTES];
  uint8_t rho[SEEDBYTES];
  uint8_t mu[CRHBYTES];
  uint8_t c[SEEDBYTES];
  uint8_t c2[SEEDBYTES];
  poly cp;
  polyvecl mat[K], z;
  polyveck t1, w1, h;
  keccak_state state;

  if(siglen != CRYPTO_BYTES)
    return -1;

  unpack_pk(rho, &t1, pk);
  if(unpack_sig(c, &z, &h, sig))
    return -1;
  if(polyvecl_chknorm(&z, GAMMA1 - BETA))
    return -1;

  /* Compute CRH(CRH(rho, t1), msg) */
  crh(mu, pk, CRYPTO_PUBLICKEYBYTES);
  shake256_init(&state);
  shake256_absorb(&state, mu, CRHBYTES);
  shake256_absorb(&state, m, mlen);
  shake256_finalize(&state);
  shake256_squeeze(mu, CRHBYTES, &state);

  /* Matrix-vector multiplication; compute Az - c2^dt1 */
  poly_challenge(&cp, c);
  polyvec_matrix_expand(mat, rho);

  polyvecl_ntt(&z);
  polyvec_matrix_pointwise_montgomery(&w1, mat, &z);

  poly_ntt(&cp);
  polyveck_shiftl(&t1);
  polyveck_ntt(&t1);
  polyveck_pointwise_poly_montgomery(&t1, &cp, &t1);

  polyveck_sub(&w1, &w1, &t1);
  polyveck_reduce(&w1);
  polyveck_invntt_tomont(&w1);

  /* Reconstruct w1 */
  polyveck_caddq(&w1);
  polyveck_use_hint(&w1, &w1, &h);
  polyveck_pack_w1(buf, &w1);

  /* Call random oracle and verify challenge */
  shake256_init(&state);
  shake256_absorb(&state, mu, CRHBYTES);
  shake256_absorb(&state, buf, K*POLYW1_PACKEDBYTES);
  shake256_finalize(&state);
  shake256_squeeze(c2, SEEDBYTES, &state);
  for(i = 0; i < SEEDBYTES; ++i)
    if(c[i] != c2[i])
      return -1;

  return 0;
}

/*************************************************
* Name:        crypto_sign_open
*
* Description: Verify signed message.
*
* Arguments:   - uint8_t *m: pointer to output message (allocated
*                            array with smlen bytes), can be equal to sm
*              - size_t *mlen: pointer to output length of message
*              - const uint8_t *sm: pointer to signed message
*              - size_t smlen: length of signed message
*              - const uint8_t *pk: pointer to bit-packed public key
*
* Returns 0 if signed message could be verified correctly and -1 otherwise
**************************************************/
int crypto_sign_open(uint8_t *m,
                     size_t *mlen,
                     const uint8_t *sm,
                     size_t smlen,
                     const uint8_t *pk)
{
  size_t i;

  if(smlen < CRYPTO_BYTES)
    goto badsig;

  *mlen = smlen - CRYPTO_BYTES;
  if(crypto_sign_verify(sm, CRYPTO_BYTES, sm + CRYPTO_BYTES, *mlen, pk))
    goto badsig;
  else {
    /* All good, copy msg, return 0 */
    for(i = 0; i < *mlen; ++i)
      m[i] = sm[CRYPTO_BYTES + i];
    return 0;
  }

badsig:
  /* Signature verification failed */
  *mlen = -1;
  for(i = 0; i < smlen; ++i)
    m[i] = 0;

  return -1;
}
